#include "eigen/conv2d_test.h"
#include <Eigen/Dense>
#include "eigen/conv2d.h"
#include "log.h"
#include "eigen/eigen_equal.h"
#include "tools/range.h"

namespace  ldl_eigen
{
const Eigen::MatrixXf input{
{1,  2,  3,  4,  5, 6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ,20, 21, 22, 23, 24, 25},
{1,  2,  3,  4,  5, 6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ,20, 21, 22, 23, 24, 25},
};

const std::vector<Eigen::MatrixXf> kernel;//
// {{
// {1,2,3},
// {4,5,6},
// {7,8,9}
// }};

const std::vector<float> bias = {1.0f};

const Eigen::MatrixXf expect_output{
{10,61,139,229,319,175,64,27,129,277,442,607,321,109,51,203,412,637,862,437,135,75,242,457,682,907,458,141,99,281,502,727,952,479,147,69,185,319,454,589,281,75,36,91,151,211,271,121,26},
{10,61,139,229,319,175,64,27,129,277,442,607,321,109,51,203,412,637,862,437,135,75,242,457,682,907,458,141,99,281,502,727,952,479,147,69,185,319,454,589,281,75,36,91,151,211,271,121,26}};

const Eigen::MatrixXf output_gradient{
{9, 8, 7, 6, 5, 4, 3,  2, 1, 9, 8, 7, 6, 5,   4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6},
{9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6, 5, 4, 3, 2, 1, 9, 8, 7, 6},
};

const Eigen::MatrixXf expect_kernel_gradient{
{1759, 1769, 1707},
{1613, 1722, 1759},
{1503, 1477, 1613}};

const Eigen::MatrixXf expect_input_gradient{
{267,267,285,258,240,168,195,267,267,285,195,150,168,195,267,258,240,195,150,168,267,285,258,240,195},
{267,267,285,258,240,168,195,267,267,285,195,150,168,195,267,258,240,195,150,168,267,285,258,240,195}};

const float expect_bias_gradient{510};

void Conv2DTest::test()
{
    const int64_t kernel_size = 3;
    const int64_t stride = 1;
    const int64_t img_height = 5;
    const int64_t img_width = 5;
    const int64_t padding = 2;

    Conv2D conv2d(img_width, img_height, kernel_size, stride, padding);
    conv2d.set_kernels(kernel);
    conv2d.set_bias(bias);

    Eigen::MatrixXf output;
    conv2d.forward(input, output);
    auto input_gradient =  conv2d.backward(output_gradient);

    LogInfo() << "input: " << input;
    LogInfo() << "img_height: " << img_height;
    LogInfo() << "img_width: " << img_width;
    LogInfo() << "kernel: " << kernel;
    LogInfo() << "stride: " << stride;
    LogInfo() << "padding: " << padding;
    LogInfo() << "bias: " << bias;
    LogInfo() << "output_gradient: " << output_gradient;
    LogInfo() << "output: " << output;
    // LogInfo() << "conv.kernel_gradient(): " << conv2d.kernel_gradient();
    // LogInfo() << "bias_gradient: " << conv2d.bias_gradient();
    LogInfo() << "conv.input_gradient(): " << conv2d.input_gradient();

    assert(EigenEqual::equal(output, expect_output, 1e-5));
    // assert(EigenEqual::equal(expect_kernel_gradient, conv2d.kernel_gradient()));
    // assert(EigenEqual::equal(expect_bias_gradient, conv2d.bias_gradient()));
    assert(EigenEqual::equal(expect_input_gradient, conv2d.input_gradient()));


    LogInfo() << "ConvTest test success.";
}
}